import pandas as pd
from pathlib import Path
import json

def process_data(filename):

    cols = ['Input','Jurisdiction','TaxType','Change','EffectiveDate']
    cols = ["date", "source", "inputSentence/Paragraph", "jurisdictionName","taxType",
            "rateChange", "effectiveDate", "jurisdictionType", "changeType", "state"]

     # this may need to chnage
    question_dict = {
        "Input": "context",
        "Jurisdiction": "Where did the change take place?",
        "TaxType": "What did change?",
        "Change": "What were the changes?",
        "EffectiveDate": "When did the change happen?"
    }

    question_dict = {
        "inputSentence/Paragraph": "context",
        "jurisdictionName": "Which location will the change take place?",
        "taxType": "What were the items or taxes that  changed?",
        "rateChange": "What were the initial and final rate changes?",
        "effectiveDate": "When will the change start?"
    }

    question_dict = {
        "inputSentence/Paragraph": "context",
        "jurisdictionName": "Which location will the change take place?",
        "taxType": "What were the items or taxes that  changed?",
        "rateChange": "What were the initial and final rate changes?",
        "effectiveDate": "When will the change start?"
    }

    """
    json format : key : { 'answers' : string [], 
                            context : string, 
                            'question' : string , 
                            'spans' = list of list[[i,j]...] 
                        }
    """


    cols_to_add = ['answers','question','spans','uuid']
    path = Path(filename)
    df = pd.read_excel(path,usecols=cols,).fillna(value = "")
    # print(df)

    df = df.rename(columns=question_dict)

    questions = [value for value in question_dict.values() if value != "context"]
    # questions = ["Where did the change take place?",
    #              "What did change?",
    #              "What were the changes?",
    #              "When did the change happen?"]


    _json_ = []

    """
    paragraph : [
    {
        "context : ""
        "qas" : [
        {
            "answers : [
                    {
                    "answer_start : "",
                    "text" : ""
                    }
                ],
                "id" : "",
                "is_impossible" : false,
                "question" : ""
        }]
        
    }]
    """
    for i in df.index:
        context = df.iloc[i]['context'].lower()
        qas = []
        for q in questions:
            text = df.iloc[i][q].lower()
            answers = map(lambda a : a.strip(),(text).split("::"))
            answer_list = []
            for answer in answers:
                answer_index = context.find(answer)
                if len(answer) > 0 and answer_index > -1:
                    answer_list.append({ "answer_start" : answer_index,
                          "text": answer
                        })
            d = {
                    "answers" : answer_list,
                    "id" : str(i)+str(questions.index(q)),
                    "is_impossible" : False,
                    "question": q
            }
            if len(d['answers']) == 0:
                d['is_impossible'] = True
            # if len(d['answers']) > 1:
            #     qas.append(d)
            # else:
            #     continue
            qas.append(d)
        # if len(qas) > 0:
        yield {
                "context" : context,
                "qas" : qas
            }
        # else:
        #     continue


def process_data_for_prediction(paragraph, output_file_dir = "/tmp/predict/"):

    """

    :param paragraph: text paragraph
    :return:
    """

    paragraph_dict = {}
    cols = ["date", "source", "inputSentence/Paragraph", "jurisdictionName", "taxType",
            "rateChange", "effectiveDate", "jurisdictionType", "changeType", "state"]

    # this may need to chnage

    question_dict = {
        "Input": "context",
        "Jurisdiction": ["Where did the change take place?","Which city or county did the change happen?"],
        "TaxType": ["What were the items or taxes that changed and what were the changes ?"],
        "Change": ["What were the changes?","What were the initial and or final changes?"],
        "EffectiveDate": ["When did the change happen?","What is the start or effective date for the change?"]
    }

    question_dict = {
        "Input": "context",
        "Jurisdiction": ["Where did the change happen?"],
        "TaxType": ["What were the items or taxes that changed?"],
        "Change": ["What were the changes?"],
        "EffectiveDate": ["What is the start or effective date for the change?"],
        "EndDate": ["What is the end date for the change?"]
    }

    """
    json format : key : { 'answers' : string [], 
                            context : string, 
                            'question' : string , 
                            'spans' = list of list[[i,j]...] 
                        }
    """

    json_output_file = output_file_dir + "test.json"

    from itertools import chain
    questions = [value for value in question_dict.values() if value != "context"]

    # flatten the list of questions
    questions = list(chain.from_iterable(questions))
    paragraphs = [paragraph]

    _json_ = []

    """
    paragraph : [
    {
        "context : ""
        "qas" : [
        {
            "answers : [
                    {
                    "answer_start : "",
                    "text" : ""
                    }
                ],
                "id" : "",
                "is_impossible" : false,
                "question" : ""
        }]

    }]
    """
    json_paragraphs = []
    for p in  paragraphs:
        context = p
        qas = []
        for q in questions:
            text = ""
            answers = map(lambda a: a.strip(), (text).split("::"))
            answer_list = []
            for answer in answers:
                answer_index = context.find(answer)
                if len(answer) > 0 and answer_index > -1:
                    answer_list.append({"answer_start": answer_index,
                                        "text": answer
                                        })
            d = {
                "answers": answer_list,
                "id": str(questions.index(q)),
                "is_impossible": False,
                "question": q
            }
            if len(d['answers']) == 0:
                d['is_impossible'] = True

            qas.append(d)

        json_paragraphs.append({
                "context" : context,
                "qas" : qas
            })
    return [
            {"paragraphs": json_paragraphs,
             "title": "Tax"
             }]
    # with open(str(json_output_file), 'w') as f:
    #
    #     json.dump({"data": [
    #         {"paragraphs": json_paragraphs,
    #          "title": "Tax"
    #          }]
    #         , "version": "v2.0"},
    #         f)
    # return  json_output_file

def token_span(sentence):
    tokens = sentence.split()
    spans = []
    for t in tokens:
        spans.append((t, sentence.find(t), sentence.find(t) + len(t)))
    return spans

if __name__ == '__main__':
    file_name = "/Users/avinash.v/Desktop/RateAnnotations_new_small.xlsx"
    test_json_output_file = "/Users/avinash.v/Desktop/sample_output_new_small.json"
    p = Path(test_json_output_file)
    json_list = (process_data(file_name))
    tuple_list = []
    count = 0
    paragraphs = []
    for d in json_list:
        paragraphs.append(d)
        # tuple_list.append((count,d))
        # print(d)
        # count += 1

    print(count)
    with open(str(p), 'w') as outfile:
        # json.dump(dict(tuple_list),outfile)
        # json.dump({"paragraphs" : paragraphs,
        #            "title" : "Tax"
        #            }, outfile)

        json.dump({"data": [
                {"paragraphs": paragraphs,
               "title": "Tax"
               }]
            , "version":"v2.0"},
        outfile)

    # print(json_list)
    # print(len(json_list))

